package com.cjx.mybatisdataauth.config;

import com.cjx.mybatisdataauth.anno.DataAuth;
import com.cjx.mybatisdataauth.mapper.DataAuthMapper;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import javax.annotation.Resource;
import java.io.StringReader;
import java.lang.reflect.Field;
import java.lang.reflect.Method;

/**
 * @author cjunxian
 * @description 数据权限拦截器
 * @date 2022/3/13
 */
//mybatis 拦截顺序Executor -> StatementHandler->ParameterHandler->ResultSetHandler
//要在分页插件之前完成sql语句的修改 应拦截Executor
@Intercepts(
        {
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        }
)
public class DataAuthInterceptor implements Interceptor {

    @Resource
    DataAuthMapper dataAuthMapper;

    /**
     * 模拟从request中取出userid
     */
    private static final String MOCK_USERID = "123";

    /**
     * 模拟从request中取出部门id
     */
    private static final String MOCK_DEPT_ID = "123";

    /**
     * 表名占位符
     */
    private static final String TABLE_NAME = "#{TABLE_NAME}";

    /**
     * 用户名占位符
     */
    private static final String USER_ID = "#{USER_ID}";

    /**
     * 部门id占位符
     */
    private static final String DEPT_ID = "#{DEPT_ID}";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        //获取拦截下的mapper
        MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
        DataAuth dataAuth = getDataAuth(mappedStatement);

        //没有注解直接放行,可根据情况补充更多的业务逻辑
        if (dataAuth == null){
            return invocation.proceed();
        }

        //获取上一个拦截器传过来的sql
        BoundSql boundSql = mappedStatement.getBoundSql(invocation.getArgs()[1]);
        String orgSql = boundSql.getSql(); //获取到当前需要被执行的SQL
        //根据指定权限点对原有sql进行修改
        String sql = modifyOrgSql(orgSql, dataAuth);

        //利用反射修改原本的sql
        Field sqlSourceField = mappedStatement.getClass().getDeclaredField("sqlSource");
        sqlSourceField.setAccessible(true);
        Object sqlSource = sqlSourceField.get(mappedStatement);
        Field boundSqlField = sqlSource.getClass().getDeclaredField("sqlSource");
        boundSqlField.setAccessible(true);
        Object boundSqlObj = boundSqlField.get(sqlSource);
        Field sqlField = boundSqlObj.getClass().getDeclaredField("sql");
        sqlField.setAccessible(true);
        sqlField.set(boundSqlObj,sql);

        return invocation.proceed();
    }

    /**
     * 通过反射获取mapper方法是否加了数据拦截注解
     */
    private DataAuth getDataAuth(MappedStatement mappedStatement) throws ClassNotFoundException {
        DataAuth dataAuth = null;
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);
        final Class<?> cls = Class.forName(className);
        final Method[] methods = cls.getMethods();
        for (Method method : methods) {
            if (method.getName().equals(methodName) && method.isAnnotationPresent(DataAuth.class)) {
                dataAuth = method.getAnnotation(DataAuth.class);
                break;
            }
        }
        return dataAuth;
    }

    /**
     * 根据权限点拼装对应sql
     * @return 拼装后的sql
     */
    private String modifyOrgSql(String orgSQql,DataAuth dataAuth) throws JSQLParserException {
        String authSql = dataAuthMapper.getAuthSql(dataAuth.name());
        if (authSql == null||authSql.isEmpty()){
            return orgSQql;
        }
        //使用CCJSqlParserUtil对sql进行解析
        //参考mybatis-plus源码{@link com.baomidou.mybatisplus.extension.plugins.inner.PaginationInnerInterceptor#autoCountSql}
        CCJSqlParserManager parserManager = new CCJSqlParserManager();
        Select select = (Select) parserManager.parse(new StringReader(orgSQql));
        PlainSelect plain = (PlainSelect) select.getSelectBody();
        Table fromItem = (Table) plain.getFromItem();
        //有别名用别名，无别名用表名，防止字段冲突报错
        String aliasTableName = fromItem.getAlias() == null ? fromItem.getName() : fromItem.getAlias().getName();
        //对authSql中的占位符进行替换
        authSql = authSql.replace(TABLE_NAME, aliasTableName).replace(USER_ID, MOCK_USERID).replace(DEPT_ID, MOCK_DEPT_ID);
        if (plain.getWhere() == null) {
            plain.setWhere(CCJSqlParserUtil.parseCondExpression(authSql));
        } else {
            plain.setWhere(new AndExpression(plain.getWhere(), CCJSqlParserUtil.parseCondExpression(authSql)));
        }
        return select.toString();
    }
}
